{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2f294fa1",
   "metadata": {},
   "source": [
    "1、导入相关依赖包"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92da6404",
   "metadata": {},
   "outputs": [
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m运行具有“Python 3.13.0”的单元格需要ipykernel包。\n",
      "\u001b[1;31m使用所需的包 <a href='command:jupyter.createPythonEnvAndSelectController'>创建 Python 环境</a>。"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ee9063d",
   "metadata": {},
   "source": [
    "2、超参数与硬件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ba97dfc",
   "metadata": {},
   "outputs": [],
   "source": [
    "BATCH_SIZE = 128\n",
    "EPOCHS     = 15\n",
    "LR         = 1e-3\n",
    "DEVICE     = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5fd3e8c4",
   "metadata": {},
   "source": [
    "3、数据加载与预处理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed58b919",
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_data(batch_size=128):\n",
    "    \"\"\"加载 MNIST 训练集与测试集\"\"\"\n",
    "    transform = transforms.Compose([\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.1307,), (0.3081,))\n",
    "    ])\n",
    "\n",
    "    train_set = datasets.MNIST(root='./data', train=True,\n",
    "                               download=True, transform=transform)\n",
    "    test_set  = datasets.MNIST(root='./data', train=False,\n",
    "                               download=True, transform=transform)\n",
    "\n",
    "    train_loader = DataLoader(train_set, batch_size=batch_size,\n",
    "                              shuffle=True,  num_workers=0)\n",
    "    test_loader  = DataLoader(test_set,  batch_size=batch_size,\n",
    "                              shuffle=False, num_workers=0)\n",
    "    return train_loader, test_loader"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4989daf3",
   "metadata": {},
   "source": [
    "4、模型构建"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df120328",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MLP(nn.Module):\n",
    "    \"\"\"两层隐藏层 MLP\"\"\"\n",
    "    def __init__(self, input_dim=784, hidden1=512, hidden2=256, num_classes=10):\n",
    "        super(MLP, self).__init__()\n",
    "        self.net = nn.Sequential(\n",
    "            nn.Flatten(),\n",
    "            nn.Linear(input_dim, hidden1),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden1, hidden2),\n",
    "            nn.ReLU(),\n",
    "            nn.Linear(hidden2, num_classes)\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.net(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac39e4c8",
   "metadata": {},
   "source": [
    "5、训练一个 epoch 的工具函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6ceaab9",
   "metadata": {},
   "outputs": [],
   "source": [
    "def train_one_epoch(model, loader, criterion, optimizer, device):\n",
    "    model.train()\n",
    "    running_loss, correct, total = 0.0, 0, 0\n",
    "\n",
    "    pbar = tqdm(loader, total=len(loader), leave=False, ncols=80)\n",
    "    for x, y in pbar:\n",
    "        x, y = x.to(device), y.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        out = model(x)\n",
    "        loss = criterion(out, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        running_loss += loss.item() * x.size(0)\n",
    "        pred = out.argmax(1)\n",
    "        correct += pred.eq(y).sum().item()\n",
    "        total += y.size(0)\n",
    "        pbar.set_description(f\"Loss={running_loss/total:.4f}\")\n",
    "\n",
    "    return running_loss / total, correct / total"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c06952a1",
   "metadata": {},
   "source": [
    "6、测试函数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d8c0dfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def test_model(model, loader, device):\n",
    "    model.eval()\n",
    "    correct, total = 0, 0\n",
    "    with torch.no_grad():\n",
    "        for x, y in loader:\n",
    "            x, y = x.to(device), y.to(device)\n",
    "            out = model(x)\n",
    "            pred = out.argmax(1)\n",
    "            correct += pred.eq(y).sum().item()\n",
    "            total += y.size(0)\n",
    "    acc = correct / total\n",
    "    print(f\"Test Accuracy = {acc:.4f}  ({correct}/{total})\")\n",
    "    return acc"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d0bde49",
   "metadata": {},
   "source": [
    "7、绘制曲线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "237fb93e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_curve(x, y1, y2, xlabel='Epoch', ylabel='Value',\n",
    "               legend1='Train Loss', legend2='Train Accuracy'):\n",
    "    plt.figure(figsize=(12, 4))\n",
    "\n",
    "    plt.subplot(1, 2, 1)\n",
    "    plt.plot(x, y1, label=legend1, color='tab:blue')\n",
    "    plt.xlabel(xlabel); plt.ylabel('Loss'); plt.title('Training Loss')\n",
    "    plt.legend(); plt.grid(True)\n",
    "\n",
    "    plt.subplot(1, 2, 2)\n",
    "    plt.plot(x, y2, label=legend2, color='tab:orange')\n",
    "    plt.xlabel(xlabel); plt.ylabel('Accuracy'); plt.title('Training Accuracy')\n",
    "    plt.legend(); plt.grid(True)\n",
    "\n",
    "    plt.tight_layout(); plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6824273b",
   "metadata": {},
   "source": [
    "8、主流程"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b7f803",
   "metadata": {},
   "outputs": [],
   "source": [
    "def main():\n",
    "    train_loader, test_loader = load_data(BATCH_SIZE)\n",
    "    model = MLP().to(DEVICE)\n",
    "    criterion = nn.CrossEntropyLoss()\n",
    "    optimizer = optim.Adam(model.parameters(), lr=LR)\n",
    "\n",
    "    loss_list, acc_list = [], []\n",
    "\n",
    "    for epoch in range(1, EPOCHS + 1):\n",
    "        loss, acc = train_one_epoch(model, train_loader, criterion, optimizer, DEVICE)\n",
    "        loss_list.append(loss); acc_list.append(acc)\n",
    "        print(f\"Epoch {epoch:02d} | Loss={loss:.4f} | Acc={acc:.4f}\")\n",
    "\n",
    "    # 测试 & 保存\n",
    "    final_acc = test_model(model, test_loader, DEVICE)\n",
    "    torch.save(model.state_dict(), \"mlp_mnist.pth\")\n",
    "    print(\"模型已保存为 mlp_mnist.pth\")\n",
    "\n",
    "    # 可视化\n",
    "    plot_curve(range(1, EPOCHS + 1), loss_list, acc_list)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01620afa",
   "metadata": {},
   "source": [
    "9、脚本入口"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ea841b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "if __name__ == \"__main__\":\n",
    "    main()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
